{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8085d744ceb233ff",
   "metadata": {},
   "source": [
    "# Vertex AI Text Embedding\n",
    "\n",
    "Imports the VertexTextEmbedding class and initializes an instance named embed_model with a specified project and location. Uses APPLICATION_DEFAULT_CREDENTIALS if no credentials is specified. The default model is `textembedding-gecko@003` in document retrival mode."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c52b0b97984c1ceb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.embeddings.vertex import VertexTextEmbedding\n",
    "\n",
    "embed_model = VertexTextEmbedding(project=\"speedy-atom-413006\", location=\"us-central1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61d58ea0808d0941",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'model_name': 'textembedding-gecko@003',\n",
       " 'embed_batch_size': 10,\n",
       " 'embed_mode': <VertexAIEmbeddingMode.RETRIEVAL_MODE: 'retrieval'>,\n",
       " 'additional_kwargs': {},\n",
       " 'class_name': 'VertexTextEmbedding'}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embed_model.dict()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c98da813ca018111",
   "metadata": {},
   "source": [
    "## Document and Query Retrival"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f6e67d1951da538",
   "metadata": {},
   "outputs": [],
   "source": [
    "embed_text_result = embed_model.get_text_embedding(\"Hello World!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f61a801502c3de8f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0.05736415088176727,\n",
       " 0.0049842665903270245,\n",
       " -0.07065856456756592,\n",
       " -0.021812528371810913,\n",
       " 0.060468606650829315]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embed_text_result[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "416ed8894817e213",
   "metadata": {},
   "outputs": [],
   "source": [
    "embed_query_result = embed_model.get_query_embedding(\"Hello World!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62510b52e204a271",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0.05158292129635811,\n",
       " -0.033334773033857346,\n",
       " -0.03221268951892853,\n",
       " -0.029282240197062492,\n",
       " 0.020004423335194588]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embed_query_result[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d10c0164acddc5d7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.7375430761259468"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from llama_index.core.base.embeddings.base import SimilarityMode\n",
    "\n",
    "embed_model.similarity(\n",
    "    embed_text_result, embed_query_result, SimilarityMode.DOT_PRODUCT\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68292f47908eabad",
   "metadata": {},
   "source": [
    "## Using the async interface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10aa2c79d07d6f77",
   "metadata": {},
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()\n",
    "\n",
    "result = await embed_model.aget_text_embedding(\"Hello World!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "596498385119ecab",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0.05733369290828705,\n",
       " 0.005178301595151424,\n",
       " -0.07033716142177582,\n",
       " -0.021963153034448624,\n",
       " 0.06050697714090347]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result[:5]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
